package com.examples.es;

import com.alibaba.fastjson.JSON;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.source.RichSourceFunction;
import org.apache.http.HttpHost;
import org.elasticsearch.action.search.SearchRequest;
import org.elasticsearch.action.search.SearchResponse;
import org.elasticsearch.client.RequestOptions;
import org.elasticsearch.client.RestClient;
import org.elasticsearch.client.RestHighLevelClient;
import org.elasticsearch.client.core.CountRequest;
import org.elasticsearch.client.core.CountResponse;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.search.SearchHit;
import org.elasticsearch.search.SearchHits;
import org.elasticsearch.search.builder.SearchSourceBuilder;
import org.elasticsearch.search.sort.SortOrder;

import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
import java.util.List;


public class IndicesQueryIncrementSource extends RichSourceFunction<HashMap<Object, Object>> {
    private static final long serialVersionUID = 33346549840611675L;
    private final String indices;
    private static RestHighLevelClient restHighLevelClient;
    private static String sortField = null;
    private static SortOrder sortOrder = SortOrder.ASC;
    private final HttpHost httpHost;
    private final Integer increment;

    IndicesQueryIncrementSource(HttpHost httpHost, String indices, Integer increment) {
        this.indices = indices;
        this.httpHost = httpHost;
        this.increment = increment;
    }
    IndicesQueryIncrementSource(HttpHost httpHost, String indices, String sortField0, SortOrder sortOrder0, Integer increment) {
        this.indices = indices;
        this.httpHost = httpHost;
        sortField = sortField0;
        sortOrder = sortOrder0;
        this.increment = increment;
    }

    @Override
    public void open(Configuration parameters) throws Exception {
        super.open(parameters);
        restHighLevelClient = new RestHighLevelClient(RestClient.builder(this.httpHost));
    }

    @Override
    public void run(SourceContext<HashMap<Object, Object>> sourceContext) throws Exception {
        int count = 0;
        int pageSize = increment;
        CountRequest countRequest = new CountRequest(new String[]{indices});

        try {
            CountResponse countResponse = restHighLevelClient.count(countRequest, RequestOptions.DEFAULT);
            count = Integer.parseInt(String.valueOf(countResponse.getCount()));
        } catch (IOException var16) {
        }

        int maxPageNum = count % pageSize != 0 ? count / pageSize + 1 : count / pageSize;

        for(int i = 0; i < maxPageNum; ++i) {
            SearchRequest request = new SearchRequest();
            request.indices(new String[]{indices});
            SearchSourceBuilder builder = (new SearchSourceBuilder()).query(QueryBuilders.matchAllQuery());
            builder.from(i == 0 ? 1 : i * pageSize);
            builder.size(pageSize);
            if(sortField != null){
                builder.sort(sortField, sortOrder);
            }
            request.source(builder);
            SearchResponse response = null;

            try {
                response = restHighLevelClient.search(request, RequestOptions.DEFAULT);
                SearchHits hits = response.getHits();
                Iterator var14 = hits.iterator();

                while(var14.hasNext()) {
                    SearchHit hit = (SearchHit)var14.next();
                    sourceContext.collect(JSON.parseObject(hit.getSourceAsString(), HashMap.class));
                }
            } catch (IOException var17) {

            }
        }

    }

    @Override
    public void cancel() {
        if(restHighLevelClient != null){
            try {
                restHighLevelClient.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

}
